/*
 * Wazuh Vulnerability Scanner - Unit Tests
 * Copyright (C) 2015, Wazuh Inc.
 * February 21, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#include "buildAllAgentListContext_test.hpp"
#include "TrampolineOsDataCache.hpp"
#include "TrampolineRemediationDataCache.hpp"
#include "TrampolineSocketDBWrapper.hpp"
#include "buildAllAgentListContext.hpp"

using TrampolineScanContext = TScanContext<TrampolineOsDataCache, GlobalData, TrampolineRemediationDataCache>;

const auto configClusterEnable = nlohmann::json::parse(R"(
    {
        "vulnerability-detection": {
            "enabled": "yes",
            "index-status": "yes",
            "feed-update-interval": "60m",
            "cti-url": "https://cti-url.com"
        },
        "managerDisabledScan": false,
        "clusterNodeName": "node_1",
        "clusterName":"clusterName",
        "clusterEnabled":true
    })");

const auto configClusterDisabled = nlohmann::json::parse(R"(
    {
        "vulnerability-detection": {
            "enabled": "yes",
            "index-status": "yes",
            "feed-update-interval": "60m",
            "cti-url": "https://cti-url.com"
        },
        "managerDisabledScan": false,
        "clusterNodeName": "node_1",
        "clusterName":"clusterName",
        "clusterEnabled":false
    })");

TEST_F(BuildAllAgentListContextTest, BuildAllAgentListContext)
{
    // Initialize policy manager
    auto policyManager = std::make_unique<PolicyManager>();
    policyManager->initialize(configClusterDisabled);

    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();
    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_)).Times(1);

    auto allAgentContext =
        std::make_shared<TBuildAllAgentListContext<TrampolineScanContext, TrampolineSocketDBWrapper>>();
    auto scanContext = std::make_shared<TrampolineScanContext>();

    allAgentContext->handleRequest(scanContext);
}

TEST_F(BuildAllAgentListContextTest, BuildAllAgentListContextWithElements)
{
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    const nlohmann::json queryResult = nlohmann::json::parse(R"(
    [
        {
            "id": 1,
            "name": "name",
            "version": "Wazuh 4.4.4",
            "ip": "192.168.0.1",
            "node_name": "node_1"
        }
    ])");

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(queryResult));

    auto allAgentContext =
        std::make_shared<TBuildAllAgentListContext<TrampolineScanContext, TrampolineSocketDBWrapper>>();

    auto scanContext = std::make_shared<TrampolineScanContext>();

    // Context is not used
    allAgentContext->handleRequest(scanContext);

    EXPECT_EQ(scanContext->m_agents.size(), 1);
}

TEST_F(BuildAllAgentListContextTest, MissingField)
{
    // Initialize policy manager
    auto policyManager = std::make_unique<PolicyManager>();
    policyManager->initialize(configClusterDisabled);

    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    nlohmann::json queryResult = nlohmann::json::parse(R"(
    [
        {
            "name": "name",
            "version": "Wazuh 4.4.4",
            "ip": "192.168.0.1",
            "node_name": "node_1"
        }
    ])");

    // Set second argument to queryResult
    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(queryResult));

    auto allAgentContext =
        std::make_shared<TBuildAllAgentListContext<TrampolineScanContext, TrampolineSocketDBWrapper>>();

    auto scanContext = std::make_shared<TrampolineScanContext>();

    allAgentContext->handleRequest(scanContext);
}

TEST_F(BuildAllAgentListContextTest, ExceptionOnDB)
{
    // Initialize policy manager
    auto policyManager = std::make_unique<PolicyManager>();
    policyManager->initialize(configClusterDisabled);

    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    const std::string agentId {"1"};

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::Throw(SocketDbWrapperException("Error on DB")));

    auto allAgentContext =
        std::make_shared<TBuildAllAgentListContext<TrampolineScanContext, TrampolineSocketDBWrapper>>();

    auto scanContext = std::make_shared<TrampolineScanContext>();

    EXPECT_THROW(allAgentContext->handleRequest(scanContext), WdbDataException);

    spSocketDBWrapperMock.reset();
}

TEST_F(BuildAllAgentListContextTest, BuildAllAgentListContextWithElementsInCluster)
{
    // Initialize policy manager
    auto policyManager = std::make_unique<PolicyManager>();
    policyManager->initialize(configClusterEnable);

    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    nlohmann::json queryResult = nlohmann::json::parse(R"(
    [
        {
            "id": 1,
            "name": "name",
            "version": "Wazuh 4.4.4",
            "ip": "192.168.0.1",
            "node_name": "node_1"
        }
    ])");

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(queryResult));

    auto allAgentContext =
        std::make_shared<TBuildAllAgentListContext<TrampolineScanContext, TrampolineSocketDBWrapper>>();

    auto scanContext = std::make_shared<TrampolineScanContext>();

    // Context is not used
    allAgentContext->handleRequest(scanContext);

    EXPECT_EQ(scanContext->m_agents.size(), 1);
}
